import numpy as np
import cv2
import onnxruntime as ort
from onnx_utils import detect, detect_pose
import os
import matplotlib.pyplot as plt


def sort_keypoints(keypoints):
    """
    Sort the keypoints by their position in the image.

    @param keypoints: 4行2列的关键点坐标，numpy数组，每行代表一个关键点的坐标(x,y)
    @return: 按左上、右上、右下、左下顺序排列的关键点坐标
    """
    index_left_up = np.argmin(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_left_down = np.argmin(np.sum(keypoints * np.array([1, -1]), axis=1))
    index_right_down = np.argmax(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_right_up = np.argmax(np.sum(keypoints * np.array([1, -1]), axis=1))
    keypoints_sorted = np.array(
        [
            keypoints[index_left_up],
            keypoints[index_right_up],
            keypoints[index_right_down],
            keypoints[index_left_down],
        ]
    )
    return keypoints_sorted


if __name__ == "__main__":
    model = ort.InferenceSession(
        r"C:\Users\Administrator\Desktop\ultralytics-action\训练结果\二维码\模型5\weights\best.onnx"
    )
    # img = cv2.imread(r"D:\desktopD\aidc_local\models\keypoint_detection\gitignored_dir\x1test\bb1.png")
    img = cv2.imread(
        r"D:\desktopD\aidc_cxftest\x1ScannerApp\app\src\main\assets\test_data\code2d_test_data\IMG_20240624_164027_3_2.png"
    )
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    out = detect_pose(
        img_rgb, model, conf_threshold=0.5, nms_threshold=0.8, need_decode=True
    )  # 对输入图像进行预处理，再使用 onnx 推理。针对 onnx 输出结果，根据阈值进行筛选和非极大抑制，输出最终结果。
    # 分数的颜色，rgb

    # img_rgb_size = img_rgb.shape[1]
    # img_rgb = cv2.resize(img_rgb, (img_rgb_size, img_rgb_size))

    print(out)
    colors1 = [255, 0, 0]
    # 矩形框的颜色
    colors2 = [0, 255, 255]
    # 关键点的颜色
    colors3 = [255, 255, 0]
    if out is not None:
        for rect in out:
            cv2.putText(
                img_rgb,
                str(round(rect[4], 2)),
                (int(rect[0] - rect[2] // 2), int(rect[1] - rect[3] // 2)),
                1,
                5,
                colors1,
                5,
            )
            for i in [5, 8, 11, 14]:
                cv2.putText(
                    img_rgb,
                    str(round(rect[i + 2], 2)),
                    (int(rect[i] - 75), int(rect[i + 1] - 20)),
                    1,
                    2,
                    colors3,
                    2,
                )
            rect = list(map(int, rect))
            cv2.rectangle(
                img_rgb,
                (rect[0] - rect[2] // 2, rect[1] - rect[3] // 2),
                (rect[0] + rect[2] // 2, rect[1] + rect[3] // 2),
                colors1,
                3,
            )
            polylines_array = np.array(
                [
                    [rect[5], rect[6]],
                    [rect[8], rect[9]],
                    [rect[11], rect[12]],
                    [rect[14], rect[15]],
                ]
            ).astype(np.int32)
            polylines_array = sort_keypoints(polylines_array)
            cv2.polylines(img_rgb, [polylines_array], True, colors2, 3)

            r = img.shape[0] // 100
            cv2.circle(img_rgb, (int(rect[5]), int(rect[6])), r, colors3, -1)
            cv2.circle(img_rgb, (int(rect[8]), int(rect[9])), r, colors3, -1)
            cv2.circle(img_rgb, (int(rect[11]), int(rect[12])), r, colors3, -1)
            cv2.circle(img_rgb, (int(rect[14]), int(rect[15])), r, colors3, -1)

    plt.imshow(img_rgb)
    plt.show()
